{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eheI-Oqc5ZKX"
      },
      "source": [
        "# Multimodal RE Tutorial\n",
        "> Tutorial author: 乔硕斐（shuofei@zju.edu.cn）\n",
        "\n",
        "In this tutorial, we use a Transformer based two-stream multimodal model——`IFAformer` to extract relations. We hope this tutorial can help you understand the process of multimodal relation extraction.\n",
        "\n",
        "This tutorial uses `Python3`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pGGj1uMN8VnK"
      },
      "source": [
        "## RE\n",
        "\n",
        "**Relation Extraction**(RE), a key task in information extraction, predicts semantic relations between pairs of entities from unstructured texts."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fkAMBMQj9Zhb"
      },
      "source": [
        "## Multimodal RE\n",
        "\n",
        "**Multimodal Relation Extraction**(MRE) in DeepKE applies a correlational visual image for each sentence, witch contains a pair of entities, to enhance the textual relation extraction."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rrln-WWWJe-w"
      },
      "source": [
        "## Dataset\n",
        "\n",
        "We use [MNRE](https://github.com/thecharm/MNRE)(Multimodal Neural Relation Extraction) dataset in this tutorial. Each piece of data contains a sentence and an image. The data formats are as follow:\n",
        "\n",
        "**Text**\n",
        "```\n",
        "{\n",
        "  'token': ['RT', '@TheHerd', ':', 'LeBron', 'is', 'in', 'their', 'head', '.', '2', '-', '0', 'Cavs', '.'],\n",
        "  'h': {\n",
        "      'name': 'LeBron',\n",
        "      'pos': [3, 4]\n",
        "     },\n",
        "  't': {\n",
        "      'name': 'Cavs',\n",
        "      'pos': [12, 13]\n",
        "     }, \n",
        "  'img_id': 'twitter_stream_2018_05_07_16_0_2_180.jpg',\n",
        "  'relation': '/per/org/member_of'\n",
        "}\n",
        "```\n",
        "**Image**\n",
        "\n",
        "![image](https://github.com/zjunlp/DeepKE/blob/main/tutorial-notebooks/re/multimodal/image/image.png?raw=1)\n",
        "\n",
        "The structure of the dataset folder `./data/` is as follow:\n",
        "\n",
        "```\n",
        ".\n",
        "├── img_detect                     # Detected objects using RCNN\n",
        "├── img_vg                       # Detected objects using visual grounding toolkit\n",
        "├── img_org                      # Original images\n",
        "├── txt                        # Text Set\n",
        "├── vg_data                      # Bounding image and img_vg\n",
        "└── ours_rel2id.json                  # Relation Set\n",
        "```\n",
        "\n",
        "We use RCNN detected objects and visual grounding objects as visual local information, where RCNN via [faster_rcnn](https://github.com/pytorch/vision/blob/main/torchvision/models/detection/faster_rcnn.py) and visual grounding via [onestage_grounding](https://github.com/zyang-ur/onestage_grounding)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LyBPnPgEMGIP"
      },
      "source": [
        "## IFAformer\n",
        "\n",
        "IFAformer is a novel dual Multimodal Transformer model with implicit feature alignment for the RE task, which utilizes the Transformer structure uniformly in the visual and textual without explicitly designing modal alignment structure.\n",
        "\n",
        "![IFAformer](https://github.com/zjunlp/DeepKE/blob/main/tutorial-notebooks/re/multimodal/image/IFAformer.png?raw=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pNqcg4kdYJO8"
      },
      "source": [
        "## Prepare evironment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "JnF_8c_QWaNR",
        "outputId": "7f1ec85d-9bf5-4fe6-8ae7-a2b621524482"
      },
      "outputs": [],
      "source": [
        "! nvidia-smi\n",
        "! pip install deepke\n",
        "! git clone https://github.com/zjunlp/DeepKE.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8jXLNAtpYYPT"
      },
      "source": [
        "Since the **torchvision** provided by Colab is the latest version, it is incompatible with the **torch** version required by deepke, so it needs to be unified. **This step can be ignored in non Colab environments.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iX2FCx_OXyRU",
        "outputId": "be64d7c1-cac9-4a47-d498-baee49845af2"
      },
      "outputs": [],
      "source": [
        "! pip install torch==1.10.0+cu102 torchvision==0.11.0+cu102 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z-0iSBR_Z3aU"
      },
      "source": [
        "## Import package"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wDNg_bVPZ8U1"
      },
      "source": [
        "The python environment provided by Colab is 3.7, which is inconsistent with the python=3.8 required by deepke. To solve the problem, it is needed to fix ***import importlib.metadata as importlib\\_metadata*** to ***import importlib\\_metadata*** in file /usr/local/lib/Python3.7/dist-packages/deepke/relation\\_extraction/multimodal/models/clip/file\\_utils.py. **For non colab environments, please use python3.8, and ignore this step.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5QXCcT9Kb1OS"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import hydra\n",
        "import torch\n",
        "import numpy as np\n",
        "import random\n",
        "from torchvision import transforms\n",
        "from hydra import utils\n",
        "import json\n",
        "import ast\n",
        "from PIL import Image\n",
        "from torch.utils.data import Dataset\n",
        "from transformers import BertTokenizer\n",
        "from torch.utils.data import DataLoader\n",
        "from deepke.relation_extraction.multimodal.models.clip.processing_clip import CLIPProcessor\n",
        "from deepke.relation_extraction.multimodal.models.IFA_model import IFAREModel\n",
        "from deepke.relation_extraction.multimodal.modules.dataset import MMREProcessor, MMREDataset\n",
        "from deepke.relation_extraction.multimodal.modules.train import Trainer\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
        "import logging\n",
        "logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',\n",
        "                    datefmt = '%m/%d/%Y %H:%M:%S',\n",
        "                    level = logging.INFO)\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "import wandb\n",
        "writer = wandb.init(project=\"DeepKE_RE_MM\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sktchqBacIpT"
      },
      "source": [
        "## Prepare dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qGaxqcV8cLR5"
      },
      "outputs": [],
      "source": [
        "! wget 120.27.214.45/Data/re/multimodal/data.tar.gz\n",
        "! tar -xzvf data.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HYbHUYA5Nqpe"
      },
      "outputs": [],
      "source": [
        "DATA_PATH = {\n",
        "        'train': 'data/txt/ours_train.txt',\n",
        "        'dev': 'data/txt/ours_val.txt',\n",
        "        'test': 'data/txt/ours_test.txt',\n",
        "        'train_auximgs': 'data/txt/mre_train_dict.pth',\n",
        "        'dev_auximgs': 'data/txt/mre_dev_dict.pth',\n",
        "        'test_auximgs': 'data/txt/mre_test_dict.pth',\n",
        "        'train_img2crop': 'data/img_detect/train/train_img2crop.pth',\n",
        "        'dev_img2crop': 'data/img_detect/val/val_img2crop.pth',\n",
        "        'test_img2crop': 'data/img_detect/test/test_img2crop.pth'}\n",
        "\n",
        "IMG_PATH = {\n",
        "    'train': 'data/img_org/train/',\n",
        "    'dev': 'data/img_org/val/',\n",
        "    'test': 'data/img_org/test'}\n",
        "\n",
        "AUX_PATH = {\n",
        "        'train': 'data/img_vg/train/crops',\n",
        "        'dev': 'data/img_vg/val/crops',\n",
        "        'test': 'data/img_vg/test/crops'}\n",
        "\n",
        "re_path = 'data/ours_rel2id.json'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TorkNH6PThuN"
      },
      "source": [
        "## Configure parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yHBaXyoKNtit"
      },
      "outputs": [],
      "source": [
        "class Config(object):\n",
        "    seed = 1234\n",
        "    \n",
        "    bert_name = \"bert-base-uncased\"\n",
        "    vit_name = \"openai/clip-vit-base-patch32\"\n",
        "    device = \"cuda\"\n",
        "    \n",
        "    num_epochs = 12\n",
        "    batch_size = 32\n",
        "    lr = 1e-5\n",
        "    warmup_ratio = 0.06\n",
        "    eval_begin_epoch = 1\n",
        "\n",
        "    max_seq = 80\n",
        "    aux_size = 128\n",
        "    rcnn_size = 64\n",
        "\n",
        "    save_path = \"checkpoints\"\n",
        "    load_path = None\n",
        "\n",
        "cfg = Config()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "34vomKiXV5SB"
      },
      "outputs": [],
      "source": [
        "def set_seed(seed=2021):\n",
        "    \"\"\"set random seed\"\"\"\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    np.random.seed(seed)\n",
        "    random.seed(seed)\n",
        "\n",
        "transform = transforms.Compose([\n",
        "    transforms.Resize(256),\n",
        "    transforms.CenterCrop(224),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
        "              std=[0.229, 0.224, 0.225])])\n",
        "\n",
        "set_seed(cfg.seed) # set seed, default is 1\n",
        "if cfg.save_path is not None:  # make save_path dir\n",
        "    if not os.path.exists(cfg.save_path):\n",
        "        os.makedirs(cfg.save_path, exist_ok=True)\n",
        "print(cfg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xSGYIyvkYJPn"
      },
      "source": [
        "## Process dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cGyx_nsCYXgX"
      },
      "outputs": [],
      "source": [
        "class MMREProcessor(object):\n",
        "    def __init__(self, data_path, re_path, args):\n",
        "        self.args = args\n",
        "        self.data_path = data_path\n",
        "        self.re_path = re_path\n",
        "        self.tokenizer = BertTokenizer.from_pretrained(args.bert_name, do_lower_case=True)\n",
        "        self.tokenizer.add_special_tokens({'additional_special_tokens':['<s>', '</s>', '<o>', '</o>']})\n",
        "\n",
        "        self.clip_processor = CLIPProcessor.from_pretrained(args.vit_name)\n",
        "        self.aux_processor = CLIPProcessor.from_pretrained(args.vit_name)\n",
        "        self.aux_processor.feature_extractor.size, self.aux_processor.feature_extractor.crop_size = args.aux_size, args.aux_size\n",
        "        self.rcnn_processor = CLIPProcessor.from_pretrained(args.vit_name)\n",
        "        self.rcnn_processor.feature_extractor.size, self.rcnn_processor.feature_extractor.crop_size = args.rcnn_size, args.rcnn_size\n",
        "\n",
        "\n",
        "    def load_from_file(self, mode=\"train\"):\n",
        "        load_file = os.path.join(self.args.cwd,self.data_path[mode])\n",
        "        logger.info(\"Loading data from {}\".format(load_file))\n",
        "        with open(load_file, \"r\", encoding=\"utf-8\") as f:\n",
        "            lines = f.readlines()\n",
        "            words, relations, heads, tails, imgids, dataid = [], [], [], [], [], []\n",
        "            for i, line in enumerate(lines):\n",
        "                line = ast.literal_eval(line)   # str to dict\n",
        "                words.append(line['token'])\n",
        "                relations.append(line['relation'])\n",
        "                heads.append(line['h']) # {name, pos}\n",
        "                tails.append(line['t'])\n",
        "                imgids.append(line['img_id'])\n",
        "                dataid.append(i)\n",
        "\n",
        "        assert len(words) == len(relations) == len(heads) == len(tails) == (len(imgids))\n",
        "\n",
        "        # 辅助图像\n",
        "        aux_imgs = None\n",
        "        # if not self.use_clip_vit:\n",
        "        aux_path = os.path.join(self.args.cwd,self.data_path[mode+\"_auximgs\"])\n",
        "        aux_imgs = torch.load(aux_path)\n",
        "        rcnn_imgs = torch.load(os.path.join(self.args.cwd,self.data_path[mode+'_img2crop']))\n",
        "        return {'words':words, 'relations':relations, 'heads':heads, 'tails':tails, 'imgids': imgids, 'dataid': dataid, 'aux_imgs':aux_imgs, \"rcnn_imgs\":rcnn_imgs}\n",
        "\n",
        "\n",
        "    def get_relation_dict(self):\n",
        "        \n",
        "        with open(os.path.join(self.args.cwd,self.re_path), 'r', encoding=\"utf-8\") as f:\n",
        "            line = f.readlines()[0]\n",
        "            re_dict = json.loads(line)\n",
        "        return re_dict\n",
        "\n",
        "    def get_rel2id(self, train_path):\n",
        "        with open(os.path.join(self.args.cwd,self.re_path), 'r', encoding=\"utf-8\") as f:\n",
        "            line = f.readlines()[0]\n",
        "            re_dict = json.loads(line)\n",
        "        re2id = {key:[] for key in re_dict.keys()}\n",
        "        with open(train_path, \"r\", encoding=\"utf-8\") as f:\n",
        "            lines = f.readlines()\n",
        "            for i, line in enumerate(lines):\n",
        "                line = ast.literal_eval(line)   # str to dict\n",
        "                assert line['relation'] in re2id\n",
        "                re2id[line['relation']].append(i)\n",
        "        return re2id\n",
        "\n",
        "\n",
        "class MMREDataset(Dataset):\n",
        "    def __init__(self, processor, transform, img_path=None, aux_img_path=None, mode=\"train\") -> None:\n",
        "        self.processor = processor\n",
        "        self.args = self.processor.args\n",
        "        self.transform = transform\n",
        "        self.max_seq = self.args.max_seq\n",
        "        self.img_path = img_path[mode]  if img_path is not None else img_path\n",
        "        self.aux_img_path = aux_img_path[mode] if aux_img_path is not None else aux_img_path\n",
        "        self.rcnn_img_path = 'data'\n",
        "        self.mode = mode\n",
        "        self.data_dict = self.processor.load_from_file(mode)\n",
        "        self.re_dict = self.processor.get_relation_dict()\n",
        "        self.tokenizer = self.processor.tokenizer\n",
        "        self.aux_size = self.args.aux_size\n",
        "        self.rcnn_size = self.args.rcnn_size\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.data_dict['words'])\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        word_list, relation, head_d, tail_d, imgid = self.data_dict['words'][idx], self.data_dict['relations'][idx], self.data_dict['heads'][idx], self.data_dict['tails'][idx], self.data_dict['imgids'][idx]\n",
        "        item_id = self.data_dict['dataid'][idx]\n",
        "        # [CLS] ... <s> head </s> ... <o> tail <o/> .. [SEP]\n",
        "        head_pos, tail_pos = head_d['pos'], tail_d['pos']\n",
        "        # insert <s> <s/> <o> <o/>\n",
        "        extend_word_list = []\n",
        "        for i in range(len(word_list)):\n",
        "            if  i == head_pos[0]:\n",
        "                extend_word_list.append('<s>')\n",
        "            if i == head_pos[1]:\n",
        "                extend_word_list.append('</s>')\n",
        "            if i == tail_pos[0]:\n",
        "                extend_word_list.append('<o>')\n",
        "            if i == tail_pos[1]:\n",
        "                extend_word_list.append('</o>')\n",
        "            extend_word_list.append(word_list[i])\n",
        "        extend_word_list = \" \".join(extend_word_list)   # list不会进行子词分词\n",
        "        encode_dict = self.tokenizer.encode_plus(text=extend_word_list, max_length=self.max_seq, truncation=True, padding='max_length')\n",
        "        input_ids, token_type_ids, attention_mask = encode_dict['input_ids'], encode_dict['token_type_ids'], encode_dict['attention_mask']\n",
        "        input_ids, token_type_ids, attention_mask = torch.tensor(input_ids), torch.tensor(token_type_ids), torch.tensor(attention_mask)\n",
        "        \n",
        "        re_label = self.re_dict[relation]   # label to id\n",
        "\n",
        "         # image process\n",
        "        if self.img_path is not None:\n",
        "            try:\n",
        "                img_path = os.path.join(os.path.join(self.args.cwd,self.img_path), imgid)\n",
        "                img_path = img_path.replace('test', 'train')\n",
        "                image = Image.open(img_path).convert('RGB')\n",
        "                image = self.processor.clip_processor(images=image, return_tensors='pt')['pixel_values'].squeeze()\n",
        "            except:\n",
        "                img_path = os.path.join(os.path.join(self.args.cwd,self.img_path), 'inf.png')\n",
        "                image = Image.open(img_path).convert('RGB')\n",
        "                image = self.processor.clip_processor(images=image, return_tensors='pt')['pixel_values'].squeeze()\n",
        "            if self.aux_img_path is not None:\n",
        "                # 辅助图像\n",
        "                aux_imgs = []\n",
        "                aux_img_paths = []\n",
        "                imgid = imgid.split(\".\")[0]\n",
        "                if item_id in self.data_dict['aux_imgs']:\n",
        "                    aux_img_paths  = self.data_dict['aux_imgs'][item_id]\n",
        "                    # print(aux_img_paths)\n",
        "                    aux_img_paths = [os.path.join(os.path.join(self.args.cwd,self.aux_img_path), path) for path in aux_img_paths]\n",
        "                # 大于3需要舍弃\n",
        "                for i in range(min(3, len(aux_img_paths))):\n",
        "                    aux_img = Image.open(aux_img_paths[i]).convert('RGB')\n",
        "                    aux_img = self.processor.aux_processor(images=aux_img, return_tensors='pt')['pixel_values'].squeeze()\n",
        "                    aux_imgs.append(aux_img)\n",
        "\n",
        "                #小于3需要加padding-0\n",
        "                for i in range(3-len(aux_imgs)):\n",
        "                    aux_imgs.append(torch.zeros((3, self.aux_size, self.aux_size))) \n",
        "\n",
        "                aux_imgs = torch.stack(aux_imgs, dim=0)\n",
        "                assert len(aux_imgs) == 3\n",
        "\n",
        "                if self.rcnn_img_path is not None:\n",
        "                    rcnn_imgs = []\n",
        "                    rcnn_img_paths = []\n",
        "                    if imgid in self.data_dict['rcnn_imgs']:\n",
        "                        rcnn_img_paths = self.data_dict['rcnn_imgs'][imgid]\n",
        "                        rcnn_img_paths = [os.path.join(os.path.join(self.args.cwd,self.rcnn_img_path), path) for path in rcnn_img_paths]\n",
        "                     # 大于3需要舍弃\n",
        "                    for i in range(min(3, len(rcnn_img_paths))):\n",
        "                        rcnn_img = Image.open(rcnn_img_paths[i]).convert('RGB')\n",
        "                        rcnn_img = self.processor.rcnn_processor(images=rcnn_img, return_tensors='pt')['pixel_values'].squeeze()\n",
        "                        rcnn_imgs.append(rcnn_img)\n",
        "                    #小于3需要加padding-0\n",
        "                    for i in range(3-len(rcnn_imgs)):\n",
        "                        rcnn_imgs.append(torch.zeros((3, self.rcnn_size, self.rcnn_size))) \n",
        "\n",
        "                    rcnn_imgs = torch.stack(rcnn_imgs, dim=0)\n",
        "                    assert len(rcnn_imgs) == 3\n",
        "                    return input_ids, token_type_ids, attention_mask, torch.tensor(re_label), image, aux_imgs, rcnn_imgs\n",
        "\n",
        "                return input_ids, token_type_ids, attention_mask, torch.tensor(re_label), image, aux_imgs\n",
        "            \n",
        "        return input_ids, token_type_ids, attention_mask, torch.tensor(re_label)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BuMhvp1FZkDL"
      },
      "outputs": [],
      "source": [
        "processor = MMREProcessor(DATA_PATH, re_path, cfg)\n",
        "train_dataset = MMREDataset(processor, transform, IMG_PATH, AUX_PATH, mode='train')\n",
        "train_dataloader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=4, pin_memory=True)\n",
        "\n",
        "dev_dataset = MMREDataset(processor, transform, IMG_PATH, AUX_PATH, mode='dev')\n",
        "dev_dataloader = DataLoader(dev_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=4, pin_memory=True)\n",
        "\n",
        "test_dataset = MMREDataset(processor, transform, IMG_PATH, AUX_PATH, mode='test')\n",
        "test_dataloader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=4, pin_memory=True)\n",
        "\n",
        "re_dict = processor.get_relation_dict()\n",
        "num_labels = len(re_dict)\n",
        "tokenizer = processor.tokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zHwUcOHKZDdj"
      },
      "source": [
        "## Model construction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ymuijbUWYwmV"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "\n",
        "import torch.nn.functional as F\n",
        "from .modeling_IFA import IFAModel\n",
        "from .clip.modeling_clip import CLIPModel\n",
        "from .clip.configuration_clip import CLIPConfig\n",
        "from transformers import BertConfig, BertModel\n",
        "\n",
        "class IFAREModel(nn.Module):\n",
        "    def __init__(self, num_labels, tokenizer, args):\n",
        "        super(IFAREModel, self).__init__()\n",
        "\n",
        "        self.args = args\n",
        "        self.vision_config = CLIPConfig.from_pretrained(self.args.vit_name).vision_config\n",
        "        self.text_config = BertConfig.from_pretrained(self.args.bert_name)\n",
        "\n",
        "        clip_model_dict = CLIPModel.from_pretrained(self.args.vit_name).vision_model.state_dict()\n",
        "        bert_model_dict = BertModel.from_pretrained(self.args.bert_name).state_dict()\n",
        "\n",
        "        print(self.vision_config)\n",
        "        print(self.text_config)\n",
        "\n",
        "        # for re\n",
        "        self.vision_config.device = args.device\n",
        "        self.model = IFAModel(self.vision_config, self.text_config)\n",
        "\n",
        "        # load:\n",
        "        vision_names, text_names = [], []\n",
        "        model_dict = self.model.state_dict()\n",
        "        for name in model_dict:\n",
        "            if 'vision' in name:\n",
        "                clip_name = name.replace('vision_', '').replace('model.', '')\n",
        "                if clip_name in clip_model_dict:\n",
        "                    vision_names.append(clip_name)\n",
        "                    model_dict[name] = clip_model_dict[clip_name]\n",
        "            elif 'text' in name:\n",
        "                text_name = name.replace('text_', '').replace('model.', '')\n",
        "                if text_name in bert_model_dict:\n",
        "                    text_names.append(text_name)\n",
        "                    model_dict[name] = bert_model_dict[text_name]\n",
        "        assert len(vision_names) == len(clip_model_dict) and len(text_names) == len(bert_model_dict), \\\n",
        "                    (len(vision_names), len(text_names), len(clip_model_dict), len(bert_model_dict))\n",
        "        self.model.load_state_dict(model_dict)\n",
        "\n",
        "        self.model.resize_token_embeddings(len(tokenizer))\n",
        "\n",
        "        self.dropout = nn.Dropout(0.5)\n",
        "        self.classifier = nn.Linear(self.text_config.hidden_size*2, num_labels)\n",
        "        self.head_start = tokenizer.convert_tokens_to_ids(\"<s>\")\n",
        "        self.tail_start = tokenizer.convert_tokens_to_ids(\"<o>\")\n",
        "        self.tokenizer = tokenizer\n",
        "\n",
        "    def forward(\n",
        "            self, \n",
        "            input_ids=None, \n",
        "            attention_mask=None, \n",
        "            token_type_ids=None, \n",
        "            labels=None, \n",
        "            images=None, \n",
        "            aux_imgs=None,\n",
        "            rcnn_imgs=None,\n",
        "    ):\n",
        "        bsz = input_ids.size(0)\n",
        "        output = self.model(input_ids=input_ids,\n",
        "                            attention_mask=attention_mask,\n",
        "                            token_type_ids=token_type_ids,\n",
        "\n",
        "                            pixel_values=images,\n",
        "                            aux_values=aux_imgs, \n",
        "                            rcnn_values=rcnn_imgs,\n",
        "                            return_dict=True,)\n",
        "\n",
        "        last_hidden_state, pooler_output = output.last_hidden_state, output.pooler_output\n",
        "        bsz, seq_len, hidden_size = last_hidden_state.shape\n",
        "        entity_hidden_state = torch.Tensor(bsz, 2*hidden_size) # batch, 2*hidden\n",
        "        for i in range(bsz):\n",
        "            head_idx = input_ids[i].eq(self.head_start).nonzero().item()\n",
        "            tail_idx = input_ids[i].eq(self.tail_start).nonzero().item()\n",
        "            head_hidden = last_hidden_state[i, head_idx, :].squeeze()\n",
        "            tail_hidden = last_hidden_state[i, tail_idx, :].squeeze()\n",
        "            entity_hidden_state[i] = torch.cat([head_hidden, tail_hidden], dim=-1)\n",
        "        entity_hidden_state = entity_hidden_state.to(self.args.device)\n",
        "        logits = self.classifier(entity_hidden_state)\n",
        "        if labels is not None:\n",
        "            loss_fn = nn.CrossEntropyLoss()\n",
        "            return loss_fn(logits, labels.view(-1)), logits\n",
        "        return logits"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uqa797pAZ1Mc"
      },
      "outputs": [],
      "source": [
        "model = IFAREModel(num_labels, tokenizer, cfg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1-p5jcJsaUGE"
      },
      "source": [
        "## Metric function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bgFqJU35adrr"
      },
      "outputs": [],
      "source": [
        "def eval_result(true_labels, pred_result, rel2id, logger, use_name=False):\n",
        "    correct = 0\n",
        "    total = len(true_labels)\n",
        "    correct_positive = 0\n",
        "    pred_positive = 0\n",
        "    gold_positive = 0\n",
        "\n",
        "    neg = -1\n",
        "    for name in ['NA', 'na', 'no_relation', 'Other', 'Others', 'none', 'None']:\n",
        "        if name in rel2id:\n",
        "            if use_name:\n",
        "                neg = name\n",
        "            else:\n",
        "                neg = rel2id[name]\n",
        "            break\n",
        "    for i in range(total):\n",
        "        if use_name:\n",
        "            golden = true_labels[i]\n",
        "        else:\n",
        "            golden = true_labels[i]\n",
        "\n",
        "        if golden == pred_result[i]:\n",
        "            correct += 1\n",
        "            if golden != neg:\n",
        "                correct_positive += 1\n",
        "        if golden != neg:\n",
        "            gold_positive += 1\n",
        "        if pred_result[i] != neg:\n",
        "            pred_positive += 1\n",
        "    acc = float(correct) / float(total)\n",
        "    try:\n",
        "        micro_p = float(correct_positive) / float(pred_positive)\n",
        "    except:\n",
        "        micro_p = 0\n",
        "    try:\n",
        "        micro_r = float(correct_positive) / float(gold_positive)\n",
        "    except:\n",
        "        micro_r = 0\n",
        "    try:\n",
        "        micro_f1 = 2 * micro_p * micro_r / (micro_p + micro_r)\n",
        "    except:\n",
        "        micro_f1 = 0\n",
        "\n",
        "    result = {'acc': acc, 'micro_p': micro_p, 'micro_r': micro_r, 'micro_f1': micro_f1}\n",
        "    logger.info('Evaluation result: {}.'.format(result))\n",
        "    return result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ev4vC9ZrZUwr"
      },
      "source": [
        "## Train the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ihxiLvmTZ7IS"
      },
      "outputs": [],
      "source": [
        "class Trainer(object):\n",
        "    def __init__(self, train_data=None, dev_data=None, test_data=None, re_dict=None, model=None, args=None, logger=None, writer=None) -> None:\n",
        "        self.train_data = train_data\n",
        "        self.dev_data = dev_data\n",
        "        self.test_data = test_data\n",
        "        self.re_dict = re_dict\n",
        "        self.model = model\n",
        "        self.logger = logger\n",
        "        self.refresh_step = 2\n",
        "        self.best_dev_metric = 0\n",
        "        self.best_test_metric = 0\n",
        "        self.best_dev_epoch = None\n",
        "        self.best_test_epoch = None\n",
        "        self.optimizer = None\n",
        "        self.writer = writer\n",
        "        \n",
        "        self.step = 0\n",
        "        self.args = args\n",
        "\n",
        "        if self.train_data is not None:\n",
        "            self.train_num_steps = len(self.train_data) * args.num_epochs\n",
        "            self.before_multimodal_train()\n",
        "        self.model.to(self.args.device)\n",
        "\n",
        "    def train(self):\n",
        "        self.step = 0\n",
        "        self.model.train()\n",
        "        self.logger.info(\"***** Running training *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.train_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Num epoch = %d\", self.args.num_epochs)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "        self.logger.info(\"  Learning rate = {}\".format(self.args.lr))\n",
        "        self.logger.info(\"  Evaluate begin = %d\", self.args.eval_begin_epoch)\n",
        "\n",
        "        if self.args.load_path is not None:  # load model from load_path\n",
        "            self.logger.info(\"Loading model from {}\".format(self.args.load_path))\n",
        "            self.model.load_state_dict(torch.load(self.args.load_path))\n",
        "            self.logger.info(\"Load model successful!\")\n",
        "\n",
        "\n",
        "        with tqdm(total=self.train_num_steps, postfix='loss:{0:<6.5f}', leave=False, dynamic_ncols=True, initial=self.step) as pbar:\n",
        "            self.pbar = pbar\n",
        "            avg_loss = 0\n",
        "            for epoch in range(1, self.args.num_epochs+1):\n",
        "                pbar.set_description_str(desc=\"Epoch {}/{}\".format(epoch, self.args.num_epochs))\n",
        "                for batch in self.train_data:\n",
        "                    self.step += 1\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)\n",
        "                    (loss, logits), labels = self._step(batch, mode=\"train\")\n",
        "                    avg_loss += loss.detach().cpu().item()\n",
        "\n",
        "                    loss.backward()\n",
        "                    self.optimizer.step()\n",
        "                    self.scheduler.step()\n",
        "                    self.optimizer.zero_grad()\n",
        "\n",
        "                    if self.step % self.refresh_step == 0:\n",
        "                        avg_loss = float(avg_loss) / self.refresh_step\n",
        "                        print_output = \"loss:{:<6.5f}\".format(avg_loss)\n",
        "                        pbar.update(self.refresh_step)\n",
        "                        pbar.set_postfix_str(print_output)\n",
        "                        self.writer.log({'avg_loss': avg_loss})\n",
        "                        avg_loss = 0\n",
        "\n",
        "                if epoch >= self.args.eval_begin_epoch:\n",
        "                    if self.dev_data:\n",
        "                        self.evaluate(epoch)   # generator to dev.\n",
        "                    if self.test_data:\n",
        "                        self.test(epoch)\n",
        "            \n",
        "            pbar.close()\n",
        "            self.pbar = None\n",
        "            self.logger.info(\"Get best dev performance at epoch {}, best dev f1 score is {}\".format(self.best_dev_epoch, self.best_dev_metric))\n",
        "            self.logger.info(\"Get best test performance at epoch {}, best test f1 score is {}\".format(self.best_test_epoch, self.best_test_metric))\n",
        "\n",
        "    def evaluate(self, epoch):\n",
        "        self.model.eval()\n",
        "        self.logger.info(\"***** Running evaluate *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.dev_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "        step = 0\n",
        "        true_labels, pred_labels = [], []\n",
        "        with torch.no_grad():\n",
        "            with tqdm(total=len(self.dev_data), leave=False, dynamic_ncols=True) as pbar:\n",
        "                pbar.set_description_str(desc=\"Dev\")\n",
        "                total_loss = 0\n",
        "                for batch in self.dev_data:\n",
        "                    step += 1\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device\n",
        "                    (loss, logits), labels = self._step(batch, mode=\"dev\")    # logits: batch, 3\n",
        "                    total_loss += loss.detach().cpu().item()\n",
        "                    \n",
        "                    preds = logits.argmax(-1)\n",
        "                    true_labels.extend(labels.view(-1).detach().cpu().tolist())\n",
        "                    pred_labels.extend(preds.view(-1).detach().cpu().tolist())\n",
        "                    pbar.update()\n",
        "                # evaluate done\n",
        "                pbar.close()\n",
        "                sk_result = classification_report(y_true=true_labels, y_pred=pred_labels, labels=list(self.re_dict.values())[1:], target_names=list(self.re_dict.keys())[1:], digits=4)\n",
        "                self.logger.info(\"%s\\n\", sk_result)\n",
        "                result = eval_result(true_labels, pred_labels, self.re_dict, self.logger)\n",
        "                acc, micro_f1 = round(result['acc']*100, 4), round(result['micro_f1']*100, 4)\n",
        "\n",
        "                self.logger.info(\"Epoch {}/{}, best dev f1: {}, best epoch: {}, current dev f1 score: {}, acc: {}.\"\\\n",
        "                            .format(epoch, self.args.num_epochs, self.best_dev_metric, self.best_dev_epoch, micro_f1, acc))\n",
        "                self.writer.log({'eva_f1': micro_f1, 'eva_accuracy': acc})\n",
        "                if micro_f1 >= self.best_dev_metric:  # this epoch get best performance\n",
        "                    self.logger.info(\"Get better performance at epoch {}\".format(epoch))\n",
        "                    self.best_dev_epoch = epoch\n",
        "                    self.best_dev_metric = micro_f1 # update best metric(f1 score)\n",
        "                    if self.args.save_path is not None:\n",
        "                        torch.save(self.model.state_dict(), self.args.save_path+\"/best_model.pth\")\n",
        "                        self.logger.info(\"Save best model at {}\".format(self.args.save_path))\n",
        "\n",
        "        self.model.train()\n",
        "\n",
        "    def test(self, epoch):\n",
        "        self.model.eval()\n",
        "        self.logger.info(\"\\n***** Running testing *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.test_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "        \n",
        "        if self.args.load_path is not None:  # load model from load_path\n",
        "            self.logger.info(\"Loading model from {}\".format(self.args.load_path))\n",
        "            self.model.load_state_dict(torch.load(self.args.load_path))\n",
        "            self.logger.info(\"Load model successful!\")\n",
        "        true_labels, pred_labels = [], []\n",
        "        with torch.no_grad():\n",
        "            with tqdm(total=len(self.test_data), leave=False, dynamic_ncols=True) as pbar:\n",
        "                pbar.set_description_str(desc=\"Testing\")\n",
        "                total_loss = 0\n",
        "                for batch in self.test_data:\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device  \n",
        "                    (loss, logits), labels = self._step(batch, mode=\"dev\")    # logits: batch, 3\n",
        "                    total_loss += loss.detach().cpu().item()\n",
        "                    \n",
        "                    preds = logits.argmax(-1)\n",
        "                    true_labels.extend(labels.view(-1).detach().cpu().tolist())\n",
        "                    pred_labels.extend(preds.view(-1).detach().cpu().tolist())\n",
        "                    \n",
        "                    pbar.update()\n",
        "                # evaluate done\n",
        "                pbar.close()\n",
        "                sk_result = classification_report(y_true=true_labels, y_pred=pred_labels, labels=list(self.re_dict.values())[1:], target_names=list(self.re_dict.keys())[1:], digits=4)\n",
        "                self.logger.info(\"%s\\n\", sk_result)\n",
        "                result = eval_result(true_labels, pred_labels, self.re_dict, self.logger)\n",
        "                acc, micro_f1 = round(result['acc']*100, 4), round(result['micro_f1']*100, 4)\n",
        "                total_loss = 0\n",
        "                self.logger.info(\"Epoch {}/{}, best test f1: {}, best epoch: {}, current test f1 score: {}, acc: {}\"\\\n",
        "                            .format(epoch, self.args.num_epochs, self.best_test_metric, self.best_test_epoch, micro_f1, acc))\n",
        "                self.writer.log({'test_f1': micro_f1, 'test_accuracy': acc})\n",
        "                if micro_f1 >= self.best_test_metric:  # this epoch get best performance\n",
        "                    self.best_test_metric = micro_f1\n",
        "                    self.best_test_epoch = epoch\n",
        "        \n",
        "        self.model.train()\n",
        "    \n",
        "    def predict(self):\n",
        "        self.model.eval()\n",
        "        self.logger.info(\"\\n***** Running predicting *****\")\n",
        "        self.logger.info(\"  Num instance = %d\", len(self.test_data)*self.args.batch_size)\n",
        "        self.logger.info(\"  Batch size = %d\", self.args.batch_size)\n",
        "        \n",
        "        if self.args.load_path is not None:  # load model from load_path\n",
        "            self.logger.info(\"Loading model from {}\".format(self.args.load_path))\n",
        "            self.model.load_state_dict(torch.load(self.args.load_path))\n",
        "            self.logger.info(\"Load model successful!\")\n",
        "        true_labels, pred_labels = [], []\n",
        "        with torch.no_grad():\n",
        "            with tqdm(total=len(self.test_data), leave=False, dynamic_ncols=True) as pbar:\n",
        "                pbar.set_description_str(desc=\"Predicting\")\n",
        "                total_loss = 0\n",
        "                for batch in self.test_data:\n",
        "                    batch = (tup.to(self.args.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device  \n",
        "                    (loss, logits), labels = self._step(batch, mode=\"dev\")    # logits: batch, 3\n",
        "                    total_loss += loss.detach().cpu().item()\n",
        "                    \n",
        "                    preds = logits.argmax(-1)\n",
        "                    true_labels.extend(labels.view(-1).detach().cpu().tolist())\n",
        "                    pred_labels.extend(preds.view(-1).detach().cpu().tolist())\n",
        "                    \n",
        "                    pbar.update()\n",
        "                # evaluate done\n",
        "                pbar.close()\n",
        "                sk_result = classification_report(y_true=true_labels, y_pred=pred_labels, labels=list(self.re_dict.values())[1:], target_names=list(self.re_dict.keys())[1:], digits=4)\n",
        "                self.logger.info(\"%s\\n\", sk_result)\n",
        "                # save predict results\n",
        "                import os\n",
        "                with open(os.path.join(self.args.cwd,'data/txt/result.txt'), 'w', encoding=\"utf-8\") as wf:\n",
        "                    wf.write(sk_result)\n",
        "                    print('Successful write!!')\n",
        "\n",
        "                \n",
        "        \n",
        "        self.model.train()\n",
        "\n",
        "    def _step(self, batch, mode=\"train\"):\n",
        "        input_ids, token_type_ids, attention_mask, labels, images, aux_imgs, rcnn_imgs = batch\n",
        "        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, labels=labels, images=images, aux_imgs=aux_imgs, rcnn_imgs=rcnn_imgs)\n",
        "        return outputs, labels\n",
        "    \n",
        "    def before_multimodal_train(self):\n",
        "        optimizer_grouped_parameters = []\n",
        "        params = {'lr':self.args.lr, 'weight_decay':1e-2}\n",
        "        params['params'] = []\n",
        "        for name, param in self.model.named_parameters():\n",
        "            if 'model' in name:\n",
        "                params['params'].append(param)\n",
        "        optimizer_grouped_parameters.append(params)\n",
        "\n",
        "        self.optimizer = optim.AdamW(optimizer_grouped_parameters, lr=self.args.lr)\n",
        "        self.scheduler = get_linear_schedule_with_warmup(optimizer=self.optimizer, \n",
        "                                 num_warmup_steps=self.args.warmup_ratio*self.train_num_steps, \n",
        "                                 num_training_steps=self.train_num_steps)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hcd-tR2waxJe"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(train_data=train_dataloader, dev_data=dev_dataloader, test_data=test_dataloader, re_dict=re_dict, model=model, args=cfg, logger=logger, writer=writer)\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TP4N4yGUa3_U"
      },
      "source": [
        "## Predict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MZcec0nUa7AF"
      },
      "outputs": [],
      "source": [
        "trainer = Trainer(train_data=None, dev_data=None, test_data=test_dataloader, re_dict=re_dict, model=model, args=cfg, logger=logger, writer=writer)\n",
        "trainer.predict()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "multimodal_re_tutorial.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
